import os, shutil

# 说明：这个test2相对于test1将保存模型更改为ModelCheckpoint保存最好的模型

root_path = "/home/python-test/py36-keras-demo01"
# 原始数据集解压目录的路径
original_dataset_dir = root_path + '/kaggle_original_data'
# 保存较小数据集的目录
base_dir = root_path + '/cats_and_dogs_small'
if not (os.path.exists(base_dir)):
    os.mkdir(base_dir)

# 分别对应划分后的训练、 验证和测试的目录
train_dir = os.path.join(base_dir, 'train')
if not (os.path.exists(train_dir)):
    os.mkdir(train_dir)
validation_dir = os.path.join(base_dir, 'validation')
if not (os.path.exists(validation_dir)):
    os.mkdir(validation_dir)
test_dir = os.path.join(base_dir, 'test')
if not (os.path.exists(test_dir)):
    os.mkdir(test_dir)

# 猫的训练图像目录
train_cats_dir = os.path.join(train_dir, 'cats')
if not (os.path.exists(train_cats_dir)):
    os.mkdir(train_cats_dir)

# 狗的训练图像目录
train_dogs_dir = os.path.join(train_dir, 'dogs')
if not (os.path.exists(train_dogs_dir)):
    os.mkdir(train_dogs_dir)

# 猫的验证图像目录
validation_cats_dir = os.path.join(validation_dir, 'cats')
if not (os.path.exists(validation_cats_dir)):
    os.mkdir(validation_cats_dir)

# 狗的验证图像目录
validation_dogs_dir = os.path.join(validation_dir, 'dogs')
if not (os.path.exists(validation_dogs_dir)):
    os.mkdir(validation_dogs_dir)

# 猫的测试图像目录
test_cats_dir = os.path.join(test_dir, 'cats')
if not (os.path.exists(test_cats_dir)):
    os.mkdir(test_cats_dir)

# 狗的测试图像目录
test_dogs_dir = os.path.join(test_dir, 'dogs')
if not (os.path.exists(test_dogs_dir)):
    os.mkdir(test_dogs_dir)

# ----------------------------------------------如果已经复制过了，就省略这一步----------------------------------------------
# # 将前 1000 张猫的图像复制到 train_cats_dir
# fnames = ['cat.{}.jpg'.format(i) for i in range(1000)]
# # print(fnames)
# for fname in fnames:
#     src = os.path.join(original_dataset_dir, fname)
#     dst = os.path.join(train_cats_dir, fname)
#     shutil.copyfile(src, dst)
#
# # 将接下来 500 张猫的图像复制到 validation_cats_dir
# fnames = ['cat.{}.jpg'.format(i) for i in range(1000, 1500)]
# for fname in fnames:
#     src = os.path.join(original_dataset_dir, fname)
#     dst = os.path.join(validation_cats_dir, fname)
#     shutil.copyfile(src, dst)
#
# # 将接下来的 500 张猫的图像复制到 test_cats_dir
# fnames = ['cat.{}.jpg'.format(i) for i in range(1500, 2000)]
# for fname in fnames:
#     src = os.path.join(original_dataset_dir, fname)
#     dst = os.path.join(test_cats_dir, fname)
#     shutil.copyfile(src, dst)
#
# # 将前 1000 张狗的图像复制到 train_dogs_dir
# fnames = ['dog.{}.jpg'.format(i) for i in range(1000)]
# for fname in fnames:
#     src = os.path.join(original_dataset_dir, fname)
#     dst = os.path.join(train_dogs_dir, fname)
#     shutil.copyfile(src, dst)
#
# # 将接下来 500 张狗的图像复制到 validation_dogs_dir
# fnames = ['dog.{}.jpg'.format(i) for i in range(1000, 1500)]
# for fname in fnames:
#     src = os.path.join(original_dataset_dir, fname)
#     dst = os.path.join(validation_dogs_dir, fname)
#     shutil.copyfile(src, dst)
#
# # 将接下来 500 张狗的图像复制到 test_dogs_dir
# fnames = ['dog.{}.jpg'.format(i) for i in range(1500, 2000)]
# for fname in fnames:
#     src = os.path.join(original_dataset_dir, fname)
#     dst = os.path.join(test_dogs_dir, fname)
#     shutil.copyfile(src, dst)

# 看看每个分组（训练 / 验证 / 测试）中分别包含多少张图像。
# print('total training cat images:', len(os.listdir(train_cats_dir)))
# print('total training dog images:', len(os.listdir(train_dogs_dir)))
# print('total validation cat images:', len(os.listdir(validation_cats_dir)))
# print('total validation dog images:', len(os.listdir(validation_dogs_dir)))
# print('total test cat images:', len(os.listdir(test_cats_dir)))
# print('total test dog images:', len(os.listdir(test_dogs_dir)))
#


# -----------------------显示几个随机增强后的训练图像---------------------------------

from keras.preprocessing.image import ImageDataGenerator
from keras.preprocessing import image
import matplotlib.pyplot as plt

datagen = ImageDataGenerator(
    rotation_range=40,
    width_shift_range=0.2,
    height_shift_range=0.2,
    shear_range=0.2,
    zoom_range=0.2,
    horizontal_flip=True,
    fill_mode='nearest')

fnames = [os.path.join(train_cats_dir, fname) for fname in os.listdir(train_cats_dir)]

img_path = fnames[3]  # 选择一张图像进行增强
img = image.load_img(img_path, target_size=(150, 150))  # 读取图像并调整大小
x = image.img_to_array(img)  # 将其转换为形状 (150, 150, 3) 的 Numpy 数组
x = x.reshape((1,) + x.shape)  # 将其形状改变为 (1, 150, 150, 3) # 要求4维
i = 0
# 生成随机变换后的图像批量。循环是无限的，因此你需要在某个时刻终止循环
for batch in datagen.flow(x, batch_size=1):
    plt.figure(i)
    imgplot = plt.imshow(image.array_to_img(batch[0]))  # array_to_img 将3维 Numpy数组转换为PIL（Python Imaging Library ）图像实例。
    i += 1
    if i % 4 == 0:
        break
plt.show(block=False)
plt.show()
